#!/usr/bin/python
# author dennis
# 2022年07月20日
import copy
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation


def loadDataSet():  # 读取数据（这里只有两个特征）
    dataMat = []
    labelMat = []
    fr = open('test.txt')
    for line in fr.readlines():
        lineArr = line.strip().split()
        dataMat.append([1.0, float(lineArr[0]), float(lineArr[1])])  # 增广特征向量
        labelMat.append(int(lineArr[2]))
    return dataMat, labelMat


def perceptron(dataMat, labelMat):
    dataMatrix = np.mat(dataMat)  # [100,3]
    classLabels = np.mat(labelMat).transpose()  # [100,1]
    m, n = dataMatrix.shape
    weights = np.zeros((n, 1))
    maxCycles = 10000
    for _ in range(maxCycles):
        for j in range(m):
            i = np.random.randint(m)
            prediction = np.sign(dataMatrix[i] * weights)
            if classLabels[i] * prediction >= 0:
                weights = weights + dataMatrix[i].transpose() * classLabels[i]
            else:
                break
    print(weights)
    return weights


# 显示数据
def plotBestFit(weights):  # 画出最终分类的图
    import matplotlib.pyplot as plt
    dataMat, labelMat = loadDataSet()
    dataArr = np.array(dataMat)
    # print(dataArr)
    n = np.shape(dataArr)[0]
    # 正样本
    xcord1 = []
    ycord1 = []
    # 负样本
    xcord2 = []
    ycord2 = []
    # 根据数据集标签进行分类
    for i in range(n):
        if int(labelMat[i]) == 1:
            # 1为正样本
            xcord1.append(dataArr[i, 1])
            ycord1.append(dataArr[i, 2])
        else:
            # 0为负样本
            xcord2.append(dataArr[i, 1])
            ycord2.append(dataArr[i, 2])
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.scatter(xcord1, ycord1, s=30, c='red', marker='s')
    ax.scatter(xcord2, ycord2, s=30, c='green')
    bnew = -weights[0] / weights[2]
    anew = -weights[1] / weights[2]
    x = np.arange(-2, 2, 0.1)
    y = lambda x: anew * x + bnew
    ax.plot(x, y(x).reshape(40, 1))
    plt.xlabel('X1')
    plt.ylabel('X2')
    plt.show()


if __name__ == '__main__':
    dataMat, labelMat = loadDataSet()
    weight = perceptron(dataMat, labelMat)
    plotBestFit(weight)
